import os
import sys
sys.path.append(".")

import pyarma as pa

from unifuncs import *

def Lawson2K1(u0, h, A, g):

    c = [0, 1/2]
    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(
        pa.expmat(c[1] * h * A) * u0 + c[1] * h * pa.expmat((1 / 2) * h * A) * Gn[0]
    )
    Gn.append(g(Un[1]))

    return pa.expmat(h * A) * u0 + h * pa.expmat((1 / 2) * h * A) * Gn[1]

def Runge2K1(u0, h, A, g):

    c = [0, 1/2]
    Un = []
    Gn = []

    Un.append(u0)
    Gn.append(g(Un[0]))

    Un.append(
        pa.expmat(c[1] * h * A) * u0 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0]])
    )
    Gn.append(g(Un[1]))

    return pa.expmat(h * A) * u0 + h * pv.phiv(h * A, [Gn[0]]) + h * pa.expmat((1 / 2) * h * A) * (Gn[1] - Gn[0])\

def Runge2K2(u0, u1, h, A, g):

    c = [0, 1/2]
    Un = []
    Gn = []

    Un.append(u1)
    Gn.append(g(Un[0]))

    Un.append(
        pa.expmat(c[1] * h * A) * u1 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0], Gn[0] - g(u0)])
    )
    Gn.append(g(Un[1]))

    return pa.expmat(h * A) * u0 + h * pv.phiv(h * A, [Gn[0], Gn[0] - g(u0)]) + h * pa.expmat((1 / 2) * h * A) * (-(3 / 2) * Gn[0] + Gn[1] + (1 / 2) * g(u0))

def Krogstad4K2(u0, u1, h, A, g):

    c = [0, 1/2, 1/2, 1]
    Un = []
    Gn = []

    Un.append(u1)
    Gn.append(g(Un[0]))

    Un.append(
        pa.expmat(c[1] * h * A) * u1 + c[1] * h * pv.phiv(c[1] * h * A, [Gn[0], (1/2) * (Gn[0] - g(u0))])
    )
    Gn.append(g(Un[1]))

    Un.append(
        pa.expmat(c[2] * h * A) * u1 + c[2] * h * pv.phiv(c[2] * h * A, [Gn[0], (1/2) * (Gn[0] - g(u0))]) + c[2] * h * (-(3/2) * Gn[0] + Gn[1] + (1/2) * g(u0))
    )
    Gn.append(g(Un[2]))

    Un.append(
        pa.expmat(c[3] * h * A) * u1 + c[3] * h * pv.phiv(c[3] * h * A, [Gn[0], Gn[0] - g(u0)]) + c[3] * h * pa.expmat(c[1] * h * A) * (-(3 / 2) * Gn[0] + Gn[2] + (1 / 2) * g(u0))
    )
    Gn.append(g(Un[3]))

    return pa.expmat(h * A) * u1 + h * pv.phiv(h * A, [Gn[0], Gn[0] - g(u0)]) + h * pa.expmat(c[1] * h * A) * (-Gn[0] + (1/3) * Gn[1] + (1/3) * Gn[2] + (1/3) * g(u0)) + h * (-(1/3) * Gn[0] + (1/6) * Gn[3] + (1/6) * g(u0))

def RBKrogstad4K2(u0, u1, h, J, F):

    return Krogstad4K2(u0, u1, h, J(u1), lambda u: F(u) - J(u1) * u)